package cn.tsa.tsp;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.math.BigInteger;
import java.net.URL;
import java.net.URLConnection;
import java.security.MessageDigest;
import java.security.SecureRandom;
import java.util.Arrays;

import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.cms.ContentInfo;
import org.bouncycastle.asn1.cryptopro.CryptoProObjectIdentifiers;
import org.bouncycastle.asn1.nist.NISTObjectIdentifiers;
import org.bouncycastle.asn1.oiw.OIWObjectIdentifiers;
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
import org.bouncycastle.asn1.x500.X500NameStyle;
import org.bouncycastle.asn1.x500.style.RFC4519Style;
import org.bouncycastle.tsp.TSPException;
import org.bouncycastle.tsp.TimeStampRequest;
import org.bouncycastle.tsp.TimeStampRequestGenerator;
import org.bouncycastle.tsp.TimeStampResponse;
import org.bouncycastle.tsp.TimeStampToken;

/**
 * Time Stamping Authority (TSA) Client [RFC 3161].
 *
 * @author Diao
 */
public class TsaClient {
	private final URL url;
	private final String username;
	private final String password;
	private final MessageDigest digest;
	X500NameStyle x500NameStyle = RFC4519Style.INSTANCE;

	/**
	 * @param url      时间戳服务接口地址
	 * @param username 时间戳服务接口用户名
	 * @param password 时间戳服务接口密码
	 * @param digest
	 */
	public TsaClient(URL url, String username, String password, MessageDigest digest) {
		this.url = url;
		this.username = username;
		this.password = password;
		this.digest = digest;
	}

	public static byte[] hexStringToByteArray(String s) {
		int len = s.length();
		byte[] data = new byte[len / 2];
		for (int i = 0; i < len; i += 2) {
			data[i / 2] = (byte) ((Character.digit(s.charAt(i), 16) << 4) + Character.digit(s.charAt(i + 1), 16));
		}
		return data;
	}

	/**
	 * @param context 待申请时间戳的电子数据
	 * @return	时间戳文件
	 * @throws IOException
	 */
	public byte[] getTimeStampToken(byte[] context) throws IOException {
		// 计算哈希值
		digest.reset();
		byte[] hash = digest.digest(context);

		SecureRandom random = new SecureRandom();
		int nonce = random.nextInt();

		// 申请时间戳
		TimeStampRequestGenerator tsaGenerator = new TimeStampRequestGenerator();
		tsaGenerator.setCertReq(true);
		ASN1ObjectIdentifier oid = getHashObjectIdentifier(digest.getAlgorithm());
		TimeStampRequest request = tsaGenerator.generate(oid, hash, BigInteger.valueOf(nonce));

		// 接收时间戳
		byte[] tsaResponse = getTSAResponse(request.getEncoded());
		
		TimeStampResponse response;
		try {
			response = new TimeStampResponse(tsaResponse);
			response.validate(request);
		} catch (TSPException e) {
			throw new IOException(e);
		}
		TimeStampToken token = response.getTimeStampToken();
		if (token == null) {
			throw new IOException("Response does not have a time stamp token");
		}

		return token.getEncoded();
	}

	/**
	 * 
	 * @param hash 待申请时间戳HASH
	 * @return 时间戳文件
	 * @throws IOException
	 */
	public byte[] getTimeStampTokenFromHash(byte[] hash) throws IOException {
		SecureRandom random = new SecureRandom();
		int nonce = random.nextInt();
		// generate TSA request
		TimeStampRequestGenerator tsaGenerator = new TimeStampRequestGenerator();
		tsaGenerator.setCertReq(true);
		ASN1ObjectIdentifier oid = getHashObjectIdentifier(digest.getAlgorithm());
		TimeStampRequest request = tsaGenerator.generate(oid, hash, BigInteger.valueOf(nonce));

		// get TSA response
		byte[] tsaResponse = getTSAResponse(request.getEncoded());

		TimeStampResponse response;
		try {
			response = new TimeStampResponse(tsaResponse);
			response.validate(request);
		} catch (TSPException e) {
			throw new IOException(e);
		}

		TimeStampToken token = response.getTimeStampToken();
		if (token == null) {
			throw new IOException("Response does not have a time stamp token");
		}

		return token.getEncoded();
	}

	/**
	 * @param request	时间戳请求
	 * @return	时间戳响应
	 * @throws IOException
	 */
	private byte[] getTSAResponse(byte[] request) throws IOException {
		URLConnection tsaConnection = url.openConnection();
		tsaConnection.setDoInput(true);
		tsaConnection.setDoOutput(true);
		tsaConnection.setUseCaches(false);
		// timeout
		tsaConnection.setConnectTimeout(8000);
		tsaConnection.setRequestProperty("Content-Type", "application/timestamp-query");
		tsaConnection.setRequestProperty("Content-Transfer-Encoding", "binary");
		if ((this.username != null) && !this.username.equals("")) {
			String userPassword = this.username + ":" + this.password;
			tsaConnection.setRequestProperty("Authorization",
					"Basic " + new String(new sun.misc.BASE64Encoder().encode(userPassword.getBytes())));
		}

		OutputStream out = tsaConnection.getOutputStream();
		out.write(request);
		out.close();
		InputStream inp = tsaConnection.getInputStream();
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		byte[] buffer = new byte[1024];
		int bytesRead = 0;
		while ((bytesRead = inp.read(buffer, 0, buffer.length)) >= 0) {
			baos.write(buffer, 0, bytesRead);
		}
		byte[] respBytes = baos.toByteArray();
		String encoding = tsaConnection.getContentEncoding();
		if (encoding != null && encoding.equalsIgnoreCase("base64")) {
			sun.misc.BASE64Decoder dec = new sun.misc.BASE64Decoder();
			respBytes = dec.decodeBuffer(new String(respBytes));
		}
		return respBytes;
	}

	/**
	 * @param token 待验证时间戳
	 * @param hash	待验证hash
	 * @return	验证结果
	 * @throws IOException 
	 * @throws TSPException 
	 */
	public boolean validateTokenTimestamp(byte[] token, byte[] hash) throws TSPException, IOException {
		TimeStampToken token2 = new TimeStampToken(ContentInfo.getInstance(token));
		byte[] messageImprintDigest = token2.getTimeStampInfo().getMessageImprintDigest();
		return Arrays.equals(messageImprintDigest, hash);
	}

	/**
	 * @param token	待解析时间戳
	 * @return	时间戳解析结果
	 * @throws IOException 
	 * @throws TSPException 
	 */
	public TSTokenInfo anayleTokenInfo(byte[] token) throws TSPException, IOException {
		return new TSTokenInfo(new TimeStampToken(ContentInfo.getInstance(token)));
	}
	
	/**
	 * @param token 时间戳文件
	 * @param path 保存路径
	 * @return
	 */
	public boolean saveTimeStampToken(byte[] token,String path) {
		boolean b = bytesAsFileWriter(path,token);
		if (b) {
			return true;
		}
		return false;
	}
	
	private static boolean bytesAsFileWriter(String path, byte[] msg) {
		OutputStream fos = null;
		try {
			File file = new File(path);
			File parent = file.getParentFile();
			if ((!parent.exists()) && (!parent.mkdirs())) {
				return false;
			}
			fos = new FileOutputStream(file);
			fos.write(msg);
			fos.flush();
			return true;
		} catch (FileNotFoundException e) {
			return false;
		} catch (IOException e) {
			return false;
		} finally {
			if (fos != null) {
				try {
					fos.close();
				} catch (IOException e) {
				}
			}
		}
	}
	/**
	 * @param algorithm	hash算法名称
	 * @return	算法oid
	 */
	private ASN1ObjectIdentifier getHashObjectIdentifier(String algorithm) {
		switch (algorithm) {
		case "MD2":
			return new ASN1ObjectIdentifier(PKCSObjectIdentifiers.md2.getId());
		case "MD5":
			return new ASN1ObjectIdentifier(PKCSObjectIdentifiers.md5.getId());
		case "SHA-1":
			return new ASN1ObjectIdentifier(OIWObjectIdentifiers.idSHA1.getId());
		case "SHA-224":
			return new ASN1ObjectIdentifier(NISTObjectIdentifiers.id_sha224.getId());
		case "SHA-256":
			return new ASN1ObjectIdentifier(NISTObjectIdentifiers.id_sha256.getId());
		case "SHA-384":
			return new ASN1ObjectIdentifier(NISTObjectIdentifiers.id_sha384.getId());
		case "SHA-512":
			return new ASN1ObjectIdentifier(NISTObjectIdentifiers.id_sha512.getId());
		case "GOST3411":
			return new ASN1ObjectIdentifier(CryptoProObjectIdentifiers.gostR3411.getId());// 1.2.643.2.2.9
		case "GOST3411-2012-256":
		case "GOST3411_2012_256":
			return new ASN1ObjectIdentifier("1.2.643.7.1.1.2.2");
		case "GOST3411-2012-512":
		case "GOST3411_2012_512":
			return new ASN1ObjectIdentifier("1.2.643.7.1.1.2.3");
		default:
			return new ASN1ObjectIdentifier(algorithm);
		}
	}

	public static byte[] inputStreamtoByteArray(InputStream input) throws IOException {
		ByteArrayOutputStream output = new ByteArrayOutputStream();
		byte[] buffer = new byte[1024 * 4];
		int n = 0;
		while (-1 != (n = input.read(buffer))) {
			output.write(buffer, 0, n);
		}
		return output.toByteArray();
	}
}